#include "specdataprocess.h"

SpecDataProcess::SpecDataProcess(int nPointsPerSignalBin, int nPointsPerNoiseBin, int nRangeBin, double sampleFreq, double wavelength, double aomShiftFreq, double searchBandwidth)
{
    this->nPointsPerSignalBin = nPointsPerSignalBin;
    this->nPointsPerNoiseBin = nPointsPerNoiseBin;
    this->nRangeBin = nRangeBin;
    this->sampleFreq = sampleFreq;
    this->wavelength = wavelength;
    this->aomShiftFreq = aomShiftFreq;
    this->searchBandwidth = searchBandwidth;

    this->freqAxis = linspace(0.0, this->sampleFreq/2.0, this->nPointsPerNoiseBin/2+1);
    this->noiseFreqAxis = this->freqAxis;
    this->signalFreqAxis = linspace(0.0, this->sampleFreq/2.0, this->nPointsPerSignalBin/2+1);

    Col<double> temp = abs(freqAxis - this->aomShiftFreq + this->searchBandwidth);
    this->startSearchIndex = temp.index_min();
    temp = abs(freqAxis - this->aomShiftFreq - this->searchBandwidth);
    this->endSearchIndex = temp.index_min();
    this->objFreqAxis = this->freqAxis(span(startSearchIndex, endSearchIndex));
//    QDBG() << "startSearchIndex = " << this->startSearchIndex;
//    QDBG() << "endSearchIndex = " << this->endSearchIndex;
//    this->objFreqAxis.print("this->objFreqAxis: ");
}

void SpecDataProcess::signalPSDProcess(quint64 *rawSpecData)
{
    double nNoiseDiv = double(this->nPointsPerNoiseBin/this->nPointsPerSignalBin);
    this->rawNoisePSD.set_size(this->nPointsPerNoiseBin);
    this->rawMirrorPSD.set_size(this->nPointsPerNoiseBin);
    this->rawSignalPSD.set_size(this->nPointsPerSignalBin, this->nRangeBin);
    int index = 0;
    for (int ii = 0; ii < this->nPointsPerNoiseBin; ii++) {
        this->rawNoisePSD(ii) = double(rawSpecData[index])/nNoiseDiv;
        index++;
    }
    for (int ii = 0; ii < this->nPointsPerNoiseBin; ii++) {
        this->rawMirrorPSD(ii) = double(rawSpecData[index])/nNoiseDiv;
        index++;
    }
    for (int jj = 0; jj < this->nRangeBin; jj++) {
        for (int ii = 0; ii < this->nPointsPerSignalBin; ii++) {
            this->rawSignalPSD(ii,jj) = double(rawSpecData[index]);
            index++;
        }
    }

    this->noisePSD.set_size(this->nPointsPerNoiseBin/2+1);
    this->mirrorPSD.set_size(this->nPointsPerNoiseBin/2+1);
    this->signalPSD.set_size(this->nPointsPerSignalBin/2+1, nRangeBin);

    this->noisePSD = 2.0 * this->rawNoisePSD(span(0, this->nPointsPerNoiseBin/2));
    this->noisePSD(0) = this->rawNoisePSD(0);
    this->noisePSD(this->nPointsPerNoiseBin/2) = this->rawNoisePSD(this->nPointsPerNoiseBin/2);

    this->mirrorPSD = 2.0 * this->rawMirrorPSD(span(0, this->nPointsPerNoiseBin/2));
    this->mirrorPSD(0) = this->rawMirrorPSD(0);
    this->mirrorPSD(this->nPointsPerNoiseBin/2) = this->rawMirrorPSD(this->nPointsPerNoiseBin/2);

    this->signalPSD = 2.0 * this->rawSignalPSD(span(0, this->nPointsPerSignalBin/2),span(0, this->nRangeBin-1));
    this->signalPSD.row(0) = this->rawSignalPSD.row(0);
    this->signalPSD.row(this->nPointsPerSignalBin/2) = this->rawSignalPSD.row(this->nPointsPerSignalBin/2);
    this->signalPSD = join_cols(this->signalPSD.row(this->nPointsPerSignalBin/2),
                                this->signalPSD(span(0,this->nPointsPerSignalBin/2-1),span::all));

    this->interpolatedSignalPSD.set_size(this->nPointsPerNoiseBin/2+1, this->nRangeBin);
    for (int kk = 0; kk < this->nRangeBin; kk++) {
        Col<double> tempYY;
        interp1(this->signalFreqAxis, this->signalPSD.col(kk), this->noiseFreqAxis, tempYY);
        this->interpolatedSignalPSD.col(kk) = tempYY;
    }

//    this->interpolatedSignalPSD.save("compare.txt",raw_ascii);

//    QDBG() << this->noisePSD.n_elem;
//    this->noisePSD.print("this->noisePSD = ");
//    QDBG() << "startSearchIndex = " << startSearchIndex;
//    QDBG() << "endSearchIndex = " << endSearchIndex;

    this->objNoisePSD = this->noisePSD(span(startSearchIndex, endSearchIndex));
    this->objMirrorPSD = this->mirrorPSD(span(startSearchIndex, endSearchIndex));
    this->objSignalPSD = this->interpolatedSignalPSD(span(startSearchIndex, endSearchIndex), span::all);

    this->spectrumWidth.set_size(this->nRangeBin);
    this->spectrumHeight.set_size(this->nRangeBin);
    this->signalSNR.set_size(this->nRangeBin);
    this->losVelocity.set_size(this->nRangeBin);

    for (int kk = 0; kk < this->nRangeBin; kk++) {
        Col<double> noiseRemovedObjPSD = this->objSignalPSD.col(kk) - this->objNoisePSD +
                mean(this->objNoisePSD);
        int idxTone = noiseRemovedObjPSD.index_max();
        this->spectrumHeight(kk) = noiseRemovedObjPSD(idxTone);

        int idxLeft = idxTone - 1;
        int idxRight = idxTone + 1;
        while (idxLeft >= 0 &&
               noiseRemovedObjPSD(idxLeft) <= noiseRemovedObjPSD(idxLeft+1)) {
            idxLeft = idxLeft - 1;
        }
        while (idxRight < noiseRemovedObjPSD.n_elem &&
               noiseRemovedObjPSD(idxRight - 1) >= noiseRemovedObjPSD(idxRight)) {
            idxRight = idxRight + 1;
        }
        idxLeft = idxLeft+1;
        idxRight = idxRight-1;

        Col<double> Nfund = noiseRemovedObjPSD;
        if (idxLeft==0 && idxRight == noiseRemovedObjPSD.n_elem - 1) {
            Nfund.shed_rows(idxLeft+1, idxRight-1);
        } else {
            Nfund.shed_rows(idxLeft, idxRight);
        }
        double meanNoisePSD = mean(Nfund);
        Col<double> noiseInObjPSD = noiseRemovedObjPSD;
        noiseInObjPSD(span(idxLeft, idxRight)) = meanNoisePSD*ones<Col<double>>(idxRight-idxLeft+1);

        Col<double> Sfund = noiseRemovedObjPSD(span(idxLeft, idxRight));
        Col<double> Ffund = this->objFreqAxis(span(idxLeft, idxRight));
        Col<uint> Sindex = find(Sfund > meanNoisePSD);
        this->losVelocity(kk) = ((dot(Sfund(Sindex), Ffund(Sindex)) /
                                  sum(Sfund(Sindex)))-this->aomShiftFreq)*this->wavelength/2;

        double fwidth = mean(diff(this->objFreqAxis));
        this->spectrumWidth(kk) = fwidth*(idxRight-idxLeft+1);

        double totalNoisePower = sum(fwidth*noiseInObjPSD);
        double totalSignalPower = sum(fwidth*(Sfund(Sindex)-meanNoisePSD));
        this->signalSNR(kk) = 10.0*qLn(totalSignalPower/totalNoisePower)/qLn(10.0);
    }

    //    this->signalSNR.print("Signal SNR:");
}

void SpecDataProcess::signalPSDProcess(quint64 *rawSpecData, Col<double> preNoisePSD)
{
    double nNoiseDiv = double(this->nPointsPerNoiseBin/this->nPointsPerSignalBin);
    this->rawNoisePSD.set_size(this->nPointsPerNoiseBin);
    this->rawMirrorPSD.set_size(this->nPointsPerNoiseBin);
    this->rawSignalPSD.set_size(this->nPointsPerSignalBin, this->nRangeBin);
    int index = 0;
    for (int ii = 0; ii < this->nPointsPerNoiseBin; ii++) {
        this->rawNoisePSD(ii) = double(rawSpecData[index])/nNoiseDiv;
        index++;
    }
    for (int ii = 0; ii < this->nPointsPerNoiseBin; ii++) {
        this->rawMirrorPSD(ii) = double(rawSpecData[index])/nNoiseDiv;
        index++;
    }
    for (int jj = 0; jj < this->nRangeBin; jj++) {
        for (int ii = 0; ii < this->nPointsPerSignalBin; ii++) {
            this->rawSignalPSD(ii,jj) = double(rawSpecData[index]);
            index++;
        }
    }

    this->noisePSD.set_size(this->nPointsPerNoiseBin/2+1);
    this->mirrorPSD.set_size(this->nPointsPerNoiseBin/2+1);
    this->signalPSD.set_size(this->nPointsPerSignalBin/2+1, nRangeBin);

    this->noisePSD = 2.0 * this->rawNoisePSD(span(0, this->nPointsPerNoiseBin/2));
    this->noisePSD(0) = this->rawNoisePSD(0);
    this->noisePSD(this->nPointsPerNoiseBin/2) = this->rawNoisePSD(this->nPointsPerNoiseBin/2);

    this->mirrorPSD = 2.0 * this->rawMirrorPSD(span(0, this->nPointsPerNoiseBin/2));
    this->mirrorPSD(0) = this->rawMirrorPSD(0);
    this->mirrorPSD(this->nPointsPerNoiseBin/2) = this->rawMirrorPSD(this->nPointsPerNoiseBin/2);

    this->signalPSD = 2.0 * this->rawSignalPSD(span(0, this->nPointsPerSignalBin/2),span(0, this->nRangeBin-1));
    this->signalPSD.row(0) = this->rawSignalPSD.row(0);
    this->signalPSD.row(this->nPointsPerSignalBin/2) = this->rawSignalPSD.row(this->nPointsPerSignalBin/2);

    this->interpolatedSignalPSD.set_size(this->nPointsPerNoiseBin/2+1, this->nRangeBin);
    for (int kk = 0; kk < this->nRangeBin; kk++) {
        Col<double> tempYY;
        interp1(this->signalFreqAxis, this->signalPSD.col(kk), this->noiseFreqAxis, tempYY);
        this->interpolatedSignalPSD.col(kk) = tempYY;
    }

//    this->interpolatedSignalPSD.save("compare.txt",raw_ascii);

//    QDBG() << this->noisePSD.n_elem;
//    this->noisePSD.print("this->noisePSD = ");
//    QDBG() << "startSearchIndex = " << startSearchIndex;
//    QDBG() << "endSearchIndex = " << endSearchIndex;

    this->objNoisePSD = this->noisePSD(span(startSearchIndex, endSearchIndex));

    this->corrValue = cor(preNoisePSD, this->objNoisePSD);
    if (corrValue.at(0) < 0.8) {
        this->objNoisePSD = preNoisePSD;
        QDBG << "Noise PSD error!";
    }

    this->objMirrorPSD = this->mirrorPSD(span(startSearchIndex, endSearchIndex));
    this->objSignalPSD = this->interpolatedSignalPSD(span(startSearchIndex, endSearchIndex), span::all);

    this->spectrumWidth.set_size(this->nRangeBin);
    this->spectrumHeight.set_size(this->nRangeBin);
    this->signalSNR.set_size(this->nRangeBin);
    this->losVelocity.set_size(this->nRangeBin);

    for (int kk = 0; kk < this->nRangeBin; kk++) {
        Col<double> noiseRemovedObjPSD = this->objSignalPSD.col(kk) - this->objNoisePSD +
                mean(this->objNoisePSD);
        int idxTone = noiseRemovedObjPSD.index_max();
        this->spectrumHeight(kk) = noiseRemovedObjPSD(idxTone);

        int idxLeft = idxTone - 1;
        int idxRight = idxTone + 1;
        while (idxLeft >= 0 &&
               noiseRemovedObjPSD(idxLeft) <= noiseRemovedObjPSD(idxLeft+1)) {
            idxLeft = idxLeft - 1;
        }
        while (idxRight < noiseRemovedObjPSD.n_elem &&
               noiseRemovedObjPSD(idxRight - 1) >= noiseRemovedObjPSD(idxRight)) {
            idxRight = idxRight + 1;
        }
        idxLeft = idxLeft+1;
        idxRight = idxRight-1;

        Col<double> Nfund = noiseRemovedObjPSD;
        if (idxLeft==0 && idxRight == noiseRemovedObjPSD.n_elem - 1) {
            Nfund.shed_rows(idxLeft+1, idxRight-1);
        } else {
            Nfund.shed_rows(idxLeft, idxRight);
        }
        double meanNoisePSD = mean(Nfund);
        Col<double> noiseInObjPSD = noiseRemovedObjPSD;
        noiseInObjPSD(span(idxLeft, idxRight)) = meanNoisePSD*ones<Col<double>>(idxRight-idxLeft+1);

        Col<double> Sfund = noiseRemovedObjPSD(span(idxLeft, idxRight));
        Col<double> Ffund = this->objFreqAxis(span(idxLeft, idxRight));
        Col<uint> Sindex = find(Sfund > meanNoisePSD);
        this->losVelocity(kk) = ((dot(Sfund(Sindex), Ffund(Sindex)) /
                                  sum(Sfund(Sindex)))-this->aomShiftFreq)*this->wavelength/2;

        double fwidth = mean(diff(this->objFreqAxis));
        this->spectrumWidth(kk) = fwidth*(idxRight-idxLeft+1);

        double totalNoisePower = sum(fwidth*noiseInObjPSD);
        double totalSignalPower = sum(fwidth*(Sfund(Sindex)-meanNoisePSD));
        this->signalSNR(kk) = 10.0*qLn(totalSignalPower/totalNoisePower)/qLn(10.0);
    }
}

double SpecDataProcess::getObjNoisePSDCorr()
{
    return this->corrValue.at(0);
}

Col<double> SpecDataProcess::getFreqAxis()
{
    return this->freqAxis;
}

Col<double> SpecDataProcess::getObjFreqAxis()
{
    return this->objFreqAxis;
}

Col<double> SpecDataProcess::getNoiseFreqAxis()
{
    return this->noiseFreqAxis;
}

Col<double> SpecDataProcess::getSignalFreqAxis()
{
    return this->signalFreqAxis;
}

Col<double> SpecDataProcess::getRawNoisePSD()
{
    return this->rawNoisePSD;
}

Col<double> SpecDataProcess::getRawMirrorPSD()
{
    return this->mirrorPSD;
}

Mat<double> SpecDataProcess::getRawSignalPSD()
{
    return this->rawSignalPSD;
}

Col<double> SpecDataProcess::getObjNoisePSD()
{
    return this->objNoisePSD;
}

Col<double> SpecDataProcess::getObjMirrorPSD()
{
    return this->mirrorPSD;
}

Mat<double> SpecDataProcess::getObjSignalPSD()
{
    return this->objSignalPSD;
}

int SpecDataProcess::getObjPSDlength()
{
    return this->objFreqAxis.n_elem*this->nRangeBin;
}

void SpecDataProcess::getObjSignalPSD(double *objSignalPSD)
{
    int index = 0;
    for (int ii = 0; ii < this->objSignalPSD.n_cols; ii++) {
        for (int jj = 0; jj < this->objSignalPSD.n_rows; jj++) {
            objSignalPSD[index] = this->objSignalPSD(jj,ii);
            index++;
        }
    }
}

void SpecDataProcess::getObjSignalNoiseRemovedPSD(double *objSignalNoiseRemovedPSD)
{
//    QDBG() << "objSignalPSD.n_cols" << this->objSignalPSD.n_cols;
//    QDBG() << "objSignalPSD.n_rows" << this->objSignalPSD.n_rows;
//    QDBG() << "this->objSignalPSD.n_elem" << this->objSignalPSD.n_elem;
    int index = 0;
    for (int ii = 0; ii < this->objSignalPSD.n_cols; ii++) {
        for (int jj = 0; jj < this->objSignalPSD.n_rows; jj++) {
            objSignalNoiseRemovedPSD[index] = this->objSignalPSD(jj,ii) - this->objNoisePSD(jj);
            index++;
        }
    }
}

int SpecDataProcess::getTotalPSDlength()
{
    return this->freqAxis.n_elem * this->nRangeBin;
}

void SpecDataProcess::getSignalPSD(double *signalPSD)
{
    int index = 0;
    for (int ii = 0; ii < this->signalPSD.n_cols; ii++) {
        for (int jj = 0; jj < this->signalPSD.n_rows; jj++) {
            signalPSD[index] = this->signalPSD(jj,ii);
            index++;
        }
    }
}

void SpecDataProcess::getSignalNoiseRemovedPSD(double *signalNoiseRemovedPSD)
{
    int index = 0;
    for (int ii = 0; ii < this->signalPSD.n_cols; ii++) {
        for (int jj = 0; jj < this->signalPSD.n_rows; jj++) {
            signalNoiseRemovedPSD[index] = this->signalPSD(jj,ii) - this->noisePSD(jj);
            index++;
        }
    }
}


Col<double> SpecDataProcess::getNoisePSD()
{
    return this->noisePSD;
}

Col<double> SpecDataProcess::getMirrorPSD()
{
    return this->mirrorPSD;
}

Mat<double> SpecDataProcess::getSignalPSD()
{
    return this->signalPSD;
}

Mat<double> SpecDataProcess::getInterpolatedSignalPSD()
{
    return this->interpolatedSignalPSD;
}

Col<double> SpecDataProcess::getSignalSNR()
{
    return this->signalSNR;
}

void SpecDataProcess::getSignalSNR(double *signalSNR)
{
    for (int i=0; i<this->signalSNR.n_elem; i++) {
        signalSNR[i] = this->signalSNR(i);
    }
}

Col<double> SpecDataProcess::getLosVelocity()
{
    return this->losVelocity;
}

void SpecDataProcess::getLosVelocity(double *losVelocity)
{
    for (int i=0; i<this->losVelocity.n_elem; i++) {
        losVelocity[i] = this->losVelocity(i);
    }
}

Col<double> SpecDataProcess::getSpectrumWidth()
{
    return this->spectrumWidth;
}

void SpecDataProcess::getSpectrumWidth(double *spectrumWidth)
{
    for (int i=0; i<this->spectrumWidth.n_elem; i++) {
        spectrumWidth[i] = this->spectrumWidth(i);
    }
}

Col<double> SpecDataProcess::getSpectrumHeight()
{
    return this->spectrumHeight;
}

void SpecDataProcess::getSpectrumHeight(double *spectrumHeight)
{
    for (int i=0; i<this->spectrumHeight.n_elem; i++) {
        spectrumHeight[i] = this->spectrumHeight(i);
    }
}
